# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/15 14:55 
@Author  : ZhangShenao 
@File    : 2.使用CacheBackedEmbeddings缓存Embedding结果.py
@Desc    : 使用CacheBackedEmbeddings缓存Embedding结果

对于相同原始内容,Embedding后的向量结果一定是相同的
因此可以增加Embedding缓存,对于相同的原始内容,直接从缓存中获取Embedding向量
可以有效减少接口的调用次数和Token的消耗,节约算力,并大大提升Embedding性能

CacheBackedEmbeddings组件,即在底层的Embedding基础上包装了缓存功能

"""
import os

import dotenv
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_openai import OpenAIEmbeddings

from cosine_similarity import cosine_similarity

# 加载环境变量
dotenv.load_dotenv()

# 创建底层Embedding模型
underlying_embeddings = OpenAIEmbeddings(
    model='text-embedding-3-small',  # 选择text-embedding-3-small模型
    dimensions=100,  # 指定Embedding生成的向量维度
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 创建CacheBackEmbedding,在底层Embedding模型的基础上,包装缓存功能
cached_embeddings = CacheBackedEmbeddings.from_bytes_store(
    underlying_embeddings=underlying_embeddings,  # 指定底层的Embedding模型
    document_embedding_cache=LocalFileStore('./cache/'),  # 指定缓存的存储位置,这里采用本地文件方式存储
    namespace=underlying_embeddings.model,  # 指定缓存的命名空间,用于隔离不同模型的Embedding结果.这里采用底层Embedding模型的名称作为命名空间
    query_embedding_cache=True,  # 开启缓存,并使用document_embedding_cache指定的存储
)

# 使用Embedding模型生成词向量列表
document_vectors = cached_embeddings.embed_documents([
    'zsa是一个喜欢打篮球的人',
    '这个喜欢打篮球的人叫zsa',
    '道阻且长，行则将至',
])

print(len(document_vectors))  # 生成3个词向量
print(len(document_vectors[0]))  # 每个词向量的维度都是100

# vector1与vector2在语义上更相似
# vector1与vector2的余弦相似度: 0.870687996049608
print(f'vector1与vector2的余弦相似度: {cosine_similarity(document_vectors[0], document_vectors[1])}')

# vector1与vector3的余弦相似度: 0.1476896129806336
print(f'vector1与vector3的余弦相似度: {cosine_similarity(document_vectors[0], document_vectors[2])}')
